/** @file admm_dev.h
 *
 *  The main ADMM driver.
 *
 *  @author Dongryeol Lee (dongryel@cc.gatech.edu)
 */

#ifndef MLPACK_ADMM_ADMM_DEV_H
#define MLPACK_ADMM_ADMM_DEV_H

#include <boost/scoped_array.hpp>
#include "mlpack/admm/admm.h"
#include "mlpack/admm/admm_coord_descent_sub_problem_solver.h"

namespace mlpack {
namespace admm {

ADMM::~ADMM() {
  if(solver_ != NULL) {
    delete solver_;
    solver_ = NULL;
  }
}

ADMM::ADMM() {
  solver_ = NULL;
  world_ = NULL;
}

void ADMM::Compute(
  const mlpack::admm::ADMMArguments &arguments_in,
  mlpack::admm::ADMMResult *result_out) {

  // Initialize the result. Minus one because one of the attributes is
  // the target.
  result_out->Init(
    arguments_in.reference_table_->num_attributes() - 1,
    arguments_in.reference_table_->num_points());

  // Auxiliary temporary variables.
  arma::vec prev_local_primal_variables;
  arma::vec current_local_primal_variables;
  arma::vec diff_local_primal_variables;
  arma::vec diff_local_dual_variables;
  arma::vec prev_local_dual_variables;
  arma::vec scaled_local_adjusted_lagrangian_variables;
  double frobenius_norm_local_primal_residuals = 0.0;
  double l2_norm_global_primal_residuals = 0.0;
  double frobenius_norm_local_dual_variables = 0.0;
  double l2_norm_global_dual_variables = 0.0;
  double frobenius_norm_local_dual_variable_change = 0.0;
  double l2_norm_global_dual_variable_change = 0.0;

  // Main iterations.
  for(int iter = 0; iter < arguments_in.max_num_iterations_ ; iter++) {

    // Update the local dual variables based on the current
    // synchronized Lagrangian variables.
    prev_local_dual_variables = result_out->local_dual_variables_;
    solver_->Solve(arguments_in, result_out);

    // Compute the local contribution of the Frobenius norm of the
    // dual variable and the Frobenius norm of the local change in the
    // dual variables.
    diff_local_dual_variables =
      result_out->local_dual_variables_ - prev_local_dual_variables;
    frobenius_norm_local_dual_variables =
      arma::norm(result_out->local_dual_variables_, "fro");
    frobenius_norm_local_dual_variable_change =
      arma::norm(diff_local_dual_variables, "fro");

    // Previous value of the primal variables backed up here.
    prev_local_primal_variables =
      result_out->local_adjusted_lagrangian_variables_ +
      result_out->synchronized_lagrangian_variables_ / arguments_in.admm_rho_;

    // Each MPI process updates its adjusted local Lagrangian
    // variables. Need to take advantage of sparsity in the local dual
    // variables $\alpha_k$.
    result_out->local_adjusted_lagrangian_variables_.zeros();
    for(int i = 0; i < arguments_in.reference_table_->num_points(); i++) {
      const core::data::alias_vec *targets = NULL;
      const core::data::alias_vec &point =
        arguments_in.reference_table_->get_instance(i, &targets);

      // The label assumed to be the last element.
      int label = static_cast<int>((*targets)[0]);
      core::math::scaled_add(
        result_out->local_dual_variables_[i] * label, point,
        &result_out->local_adjusted_lagrangian_variables_);
    }
    result_out->local_adjusted_lagrangian_variables_ -=
      result_out->synchronized_lagrangian_variables_ / arguments_in.admm_rho_;
    scaled_local_adjusted_lagrangian_variables =
      result_out->local_adjusted_lagrangian_variables_ *
      (- arguments_in.admm_rho_ /
       static_cast<double>(world_->size() + arguments_in.admm_rho_));

    // Synchronize the Lagrangian variables.
    boost::mpi::all_reduce(
      *world_, scaled_local_adjusted_lagrangian_variables.memptr(),
      scaled_local_adjusted_lagrangian_variables.n_elem,
      result_out->synchronized_lagrangian_variables_.memptr(),
      std::plus<double>());

    // The current primal variables recomputed here.
    current_local_primal_variables =
      result_out->local_adjusted_lagrangian_variables_ +
      result_out->synchronized_lagrangian_variables_ / arguments_in.admm_rho_;

    // Compute the local primal squared residuals.
    diff_local_primal_variables =
      current_local_primal_variables - prev_local_primal_variables;
    frobenius_norm_local_primal_residuals =
      arma::norm(diff_local_primal_variables, "fro");

    // Compute the squared norm difference in the primal variables and
    // eventually the L2 norm of the primal variable residual
    boost::mpi::all_reduce(
      *world_, frobenius_norm_local_primal_residuals,
      l2_norm_global_primal_residuals, std::plus<double>());
    l2_norm_global_primal_residuals = sqrt(l2_norm_global_primal_residuals);

    // Compute the L2 norm of the dual variable change and the dual variable.
    boost::mpi::all_reduce(
      *world_, frobenius_norm_local_dual_variables,
      l2_norm_global_dual_variables,
      std::plus<double>());
    l2_norm_global_dual_variables = sqrt(l2_norm_global_dual_variables);
    boost::mpi::all_reduce(
      * world_, frobenius_norm_local_dual_variable_change,
      l2_norm_global_dual_variable_change, std::plus<double>()) ;
    l2_norm_global_dual_variable_change =
      sqrt(l2_norm_global_dual_variable_change);

    // Compute primal/dual feasibility tolerances.
    double primal_residual_threshold =
      sqrt(
        arguments_in.reference_table_->average_num_points() * world_->size()) *
      arguments_in.absolute_tolerance_ +
      arguments_in.relative_tolerance_ *
      std::max<double>(
        l2_norm_global_primal_residuals,
        sqrt(world_->size()) * l2_norm_global_dual_variables);
    double dual_residual_threshold =
      sqrt(
        arguments_in.reference_table_->average_num_points() * world_->size()) *
      arguments_in.absolute_tolerance_ +
      arguments_in.relative_tolerance_ *
      arma::norm(result_out->synchronized_lagrangian_variables_ , 2) ;

    // Check convergence.
    //printf("Looking at %g vs %g and %g vs %g\n",
    //     l2_norm_global_primal_residuals, primal_residual_threshold,
    //     l2_norm_global_dual_variable_change, dual_residual_threshold);
    if(l2_norm_global_primal_residuals <= primal_residual_threshold &&
        l2_norm_global_dual_variable_change <= dual_residual_threshold) {
      break;
    }
  } // Main loop end.

  // Compute the final primal variables.
  current_local_primal_variables = - result_out->synchronized_lagrangian_variables_;
  result_out->primal_variables_ = current_local_primal_variables;
}

void ADMM::Init(
  boost::mpi::communicator &world_in,
  mlpack::admm::ADMMArguments &arguments_in) {

  // Set the communicator.
  world_ = &world_in;

  // This case implies that the query equals the reference.
  if(arguments_in.query_table_ == NULL) {
    arguments_in.query_table_ = arguments_in.reference_table_;
  }

  // Initialize the solver. By default, we are using the coordinate
  // descent solver.
  solver_ = new mlpack::admm::ADMMCoordDescentSubProblemSolver();
}
}
}

#endif
